package com.qupeng.demo.kafka.kafkaspringbootproducer.admin;

import com.qupeng.demo.kafka.kafkaspringbootproducer.admin.model.ConsumerGroup;
import com.qupeng.demo.kafka.kafkaspringbootproducer.admin.model.ConsumerGroupDescribe;
import com.qupeng.demo.kafka.kafkaspringbootproducer.admin.model.ConsumerGroupOffsets;
import com.qupeng.demo.kafka.kafkaspringbootproducer.config.KafkaProperties;
import com.qupeng.demo.kafka.kafkaspringbootproducer.model.CommonResponse;
import org.apache.kafka.clients.admin.*;
import org.apache.kafka.clients.consumer.OffsetAndMetadata;
import org.apache.kafka.common.TopicPartition;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.BeanUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.kafka.core.KafkaAdmin;
import org.springframework.web.bind.annotation.*;

import java.util.*;
import java.util.concurrent.ExecutionException;
import java.util.stream.Collectors;

@RestController
@RequestMapping("/admin/sync/consumer-groups")
public class ConsumerGroupSyncAdminController {

    private static final Logger LOGGER = LoggerFactory.getLogger(ConsumerGroupSyncAdminController.class);

    @Autowired
    private KafkaAdmin kafkaAdmin;

    @Autowired
    private AdminClient adminClient;

    @Autowired
    private KafkaProperties kafkaProperties;

    @GetMapping("")
    public CommonResponse<List<ConsumerGroup>> consumerGroups() {
        try {
            ListConsumerGroupsResult result = adminClient.listConsumerGroups();
            List<ConsumerGroup> consumerGroups = result.valid().get().stream().map(consumerGroupListing -> {
                ConsumerGroup consumerGroup = new ConsumerGroup();
                BeanUtils.copyProperties(consumerGroupListing, consumerGroup);
                return consumerGroup;
            }).collect(Collectors.toList());
            Collection<Throwable> errors = result.errors().get();
            if (!errors.isEmpty()) {
                CommonResponse<List<ConsumerGroup>> errorResponse = new CommonResponse<>("admin-consumer-group-001", "", "kafka-admin");
                errors.forEach(throwable -> {
                    errorResponse.add(new CommonResponse<>("admin-consumer-group-001", throwable.getMessage(), "kafka-admin"));
                });
                return errorResponse;
            } else {
                return new CommonResponse<>("kafka-admin", consumerGroups);
            }
        } catch (InterruptedException | ExecutionException ex) {
            return new CommonResponse<>("admin-consumer-group-001", ex.getMessage(), "kafka-admin");
        }
    }

    @GetMapping("/{name}")
    public CommonResponse<ConsumerGroupDescribe> consumerGroupDescription(@PathVariable String name) {
        try {
            ConsumerGroupDescription description = adminClient.describeConsumerGroups(Collections.singletonList(name)).describedGroups().get(name).get();
            ConsumerGroupDescribe consumerGroupDescribe = new ConsumerGroupDescribe();
            BeanUtils.copyProperties(description, consumerGroupDescribe);
            BeanUtils.copyProperties(description.coordinator(), consumerGroupDescribe.getCoordinator());
            return new CommonResponse<>("kafka-admin", consumerGroupDescribe);
        }  catch (InterruptedException | ExecutionException ex) {
            return new CommonResponse<>("admin-consumer-group-002", ex.getMessage(), "kafka-admin");
        }
    }

    @GetMapping("/{name}/offsets")
    public CommonResponse<Collection<ConsumerGroupOffsets>> consumerGroupOffsets(@PathVariable String name) {
        try {
            Map<TopicPartition, OffsetAndMetadata> offsets = adminClient.listConsumerGroupOffsets(name).partitionsToOffsetAndMetadata().get();
            Map<TopicPartition, OffsetSpec> requestLatestOffsets = new HashMap<>();
            for (TopicPartition tp : offsets.keySet()) {
                requestLatestOffsets.put(tp, OffsetSpec.latest());
            }
            Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> latestOffsets = adminClient.listOffsets(requestLatestOffsets).all().get();
            Collection<ConsumerGroupOffsets> consumerGroupOffsets = new ArrayList<>();
            for (Map.Entry<TopicPartition, OffsetAndMetadata> entry : offsets.entrySet()) {
                String topic = entry.getKey().topic();
                int partition = entry.getKey().partition();
                long committedOffset = entry.getValue().offset();
                long latestOffset = latestOffsets.get(entry.getKey()).offset();
                consumerGroupOffsets.add(new ConsumerGroupOffsets(name, topic, partition, committedOffset, latestOffset));
            }
            return new CommonResponse<>("kafka-admin", consumerGroupOffsets);
        }  catch (InterruptedException | ExecutionException ex) {
            return new CommonResponse<>("admin-consumer-group-003", ex.getMessage(), "kafka-admin");
        }
    }

    @PutMapping("/{name}/offsets")
    public CommonResponse<Collection<ConsumerGroupOffsets>> updateConsumerGroupOffsets(@PathVariable String name) {
        try {
            Map<TopicPartition, OffsetAndMetadata> offsets = adminClient.listConsumerGroupOffsets(name).partitionsToOffsetAndMetadata().get();
            Map<TopicPartition, OffsetSpec> requestEarliestOffsets = new HashMap<>();
            for (TopicPartition tp : offsets.keySet()) {
                requestEarliestOffsets.put(tp, OffsetSpec.earliest());
            }
            Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> earliestOffsets = adminClient.listOffsets(requestEarliestOffsets).all().get();
            Map<TopicPartition, OffsetAndMetadata> resetOffsets = new HashMap<>();
            for (Map.Entry<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> entry : earliestOffsets.entrySet()) {
                resetOffsets.put(entry.getKey(), new OffsetAndMetadata(entry.getValue().offset()));
            }

            adminClient.alterConsumerGroupOffsets(name, resetOffsets).all().get();

            Collection<ConsumerGroupOffsets> consumerGroupOffsets = new ArrayList<>();
            for (Map.Entry<TopicPartition, OffsetAndMetadata> entry : offsets.entrySet()) {
                String topic = entry.getKey().topic();
                int partition = entry.getKey().partition();
                long committedOffset = entry.getValue().offset();
                long latestOffset = earliestOffsets.get(entry.getKey()).offset();
                consumerGroupOffsets.add(new ConsumerGroupOffsets(name, topic, partition, committedOffset, latestOffset));
            }
            return new CommonResponse<>("kafka-admin", consumerGroupOffsets);
        }  catch (InterruptedException | ExecutionException ex) {
            return new CommonResponse<>("admin-consumer-group-003", ex.getMessage(), "kafka-admin");
        }
    }
}
